import os
import numpy as np
import torch
from torch.utils.data import DataLoader
from torchvision import transforms
import matplotlib.pyplot as plt
from utils import load_test_set
from data_proc.dataset import SegmentationDataset
from medpy.metric import dc,hd,jc,precision,recall
from utils import onehot_to_single_channel

def evaluate(cfg, batch_size=1):
    model_path = cfg['MODEL_PATH']
    best_model_path = cfg['BEST_MODEL_PATH']
    split_datasets_file_idx = cfg['SPLIT_FILE_IDX']
    pred_folder = cfg['PRED_PLOT_PATH']
    roc_curve_path = os.path.join(cfg['OUTPUT_PATH'], 'roc_curve.png')

    test_imgs, test_masks = load_test_set(split_datasets_file_idx)

    transform = transforms.Compose([
        transforms.Resize(cfg['IMAGE_SIZE']),
        transforms.ToTensor()
    ])

    test_set = SegmentationDataset(test_imgs, test_masks, transform)
    print(f'Test images: {len(test_set)}\n')

    test_loader = DataLoader(test_set, batch_size=batch_size, **cfg['KWARGS'])

    # Check if CUDA is available and move the model to GPU if so
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print(f'Using device: {device}')

    model = torch.load(model_path)
    model.to(device)  # Move model to GPU
    model.eval()


    if not os.path.exists(pred_folder):
        os.makedirs(pred_folder)
    all_metrics_results = []
    with torch.no_grad():
        for i, (inputs, targets) in enumerate(test_loader):
            inputs, targets = inputs.to(device), targets.to(device)  # Move data to GPU
            outputs = model(inputs)
            outputs = outputs.sigmoid()  # Assuming outputs are logits; apply sigmoid to convert to probabilities
            preds = outputs > 0.5  # Convert probabilities to binary output

            # Move predictions and targets to CPU for further processing and metric calculations
            preds = preds.cpu().int()  # Convert to integer for metrics calculation
            targets = targets.cpu().int()
            preds = preds.numpy().squeeze(0)
            targets = targets.numpy().squeeze(0)
            inputs = inputs.cpu()
            input_img = inputs.numpy().squeeze(0).transpose((1, 2, 0))
            # 创建一个图形和三个子图
            fig, axes = plt.subplots(1, 5, figsize=(15, 5))
            # 设置标题
            axes[0].set_title('Channel 0: Pets')
            axes[1].set_title('Channel 1: Background')
            axes[2].set_title('Channel 2: Edges')
            axes[3].set_title('True Mask')
            axes[4].set_title('Image')

            # 显示每个通道
            for j in range(3):
                axes[j].imshow(preds[j], cmap='gray')  # Display prediction channels
                axes[j].axis('off')  # Turn off axis
            targets_one_chn = onehot_to_single_channel(targets)
            axes[3].imshow(targets_one_chn, cmap='gray')  # Display true mask
            axes[3].axis('off')  # Turn off axis
            axes[4].imshow(input_img, cmap=None)  # Display true mask
            axes[4].axis('off')  # Turn off axis

            plt_path = os.path.join(pred_folder, f'pred_result_{i}.png')  # Save each plot with an index
            plt.savefig(plt_path)
            plt.close(fig)  # Close the plot to free up memory

            # 初始化存储结果
            metrics_results = {
                "Dice": {},
                "Hausdorff": {},
                "Jaccard": {},
                "Precision": {},
                "Recall": {}
            }
            # 类别映射
            class_labels = {0: 'Pets', 1: 'Background', 2: 'Edges'}

            for class_id, class_name in class_labels.items():
                # 为当前类别创建二元掩码
                pred_mask = preds[class_id].astype(np.bool_)
                true_mask = (targets_one_chn == class_id).astype(np.bool_)

                # 确保掩码形状相同，这对 MedPy 库的函数很重要
                if true_mask.shape != pred_mask.shape:
                    raise ValueError("Shape mismatch: the shape of pred_mask and true_mask must be the same.")

                # 计算指标
                try:
                    metrics_results["Dice"][class_name] = dc(pred_mask, true_mask)
                    metrics_results["Jaccard"][class_name] = jc(pred_mask, true_mask)
                    metrics_results["Hausdorff"][class_name] = hd(pred_mask, true_mask)
                    metrics_results["Precision"][class_name] = precision(pred_mask, true_mask)
                    metrics_results["Recall"][class_name] = recall(pred_mask, true_mask)
                except Exception as e:
                    print(f"Error calculating metrics for {class_name}: {e}")

            all_metrics_results.append(metrics_results)
        # Calculate average over all images
    final_results = {}
    dice_scores = []
    jaccard_scores = []
    hd_scores = []
    precision_scores = []
    recall_scores = []
    roc_aucs = []
    average_precisions = []
    # 遍历 all_metrics_results 来收集和计算每个指标的平均值
    for metrics in all_metrics_results:
        # 计算 dice 平均值
        dice_avg = np.mean(list(metrics['Dice'].values()))
        dice_scores.append(dice_avg)

        # 计算 hd 平均值
        hd_avg = np.mean(list(metrics['Hausdorff'].values()))
        hd_scores.append(hd_avg)

        # 计算 jaccard 平均值
        jc_avg = np.mean(list(metrics['Jaccard'].values()))
        jaccard_scores.append(jc_avg)

        # 计算 precision 平均值
        precision_avg = np.mean(list(metrics['Precision'].values()))
        precision_scores.append(precision_avg)

        # 计算 recall 平均值
        recall_avg = np.mean(list(metrics['Recall'].values()))
        recall_scores.append(recall_avg)

    # 打印最终的平均值
    print("Average Dice Score across all images and categories:", np.mean(dice_scores))
    print("Average HD Score across all images and categories:", np.mean(hd_scores))
    print("Average Jaccard Score across all images and categories:", np.mean(jaccard_scores))
    print("Average Precision Score across all images and categories:", np.mean(precision_scores))
    print("Average Recall Score across all images and categories:", np.mean(recall_scores))


